import torch
from torch.utils.data import Dataset, DataLoader
import os
import numpy as np

class PFNNTextDataset(Dataset):
    def __init__(self,
                 base_path = 'dataset'):
        self.create_traing_statistics()
        X = np.loadtxt(os.path.join(base_path, 'Input.txt')).astype(np.float32)
        Y = np.loadtxt(os.path.join(base_path, 'Output.txt')).astype(np.float32)
        P = np.expand_dims(X[..., -1], axis=1).astype(np.float32)
        X = np.delete(X, -1, axis=1)
        print(f'Input shape: {X.shape}')
        print(f'Output shape: {Y.shape}')
        print(f'Phase shape: {P.shape}')

        Xmean, Xstd = X.mean(axis=0), X.std(axis=0)
        Ymean, Ystd = Y.mean(axis=0), Y.std(axis=0)
        Xstd[Xstd == 0] = 1; Ystd[Ystd == 0] = 1
        X = (X - Xmean) / Xstd; Y = (Y - Ymean) / Ystd
        Xmean.tofile(self.Xmean_path)
        Ymean.tofile(self.Ymean_path)
        Xstd.tofile(self.Xstd_path)
        Ystd.tofile(self.Ystd_path)
        X = np.concatenate((X, P), axis=1)
        self.X = X; self.Y = Y; self.P = P
        self.length = len(X)

    def create_traing_statistics(self):
        statistics_path = os.path.join('training', 'nn')
        os.makedirs(statistics_path, exist_ok=True)
        self.Xmean_path = os.path.join(statistics_path, 'Xmean.bin')
        self.Xstd_path = os.path.join(statistics_path, 'Xstd.bin')
        self.Ymean_path = os.path.join(statistics_path, 'Ymean.bin')
        self.Ystd_path = os.path.join(statistics_path, 'Ystd.bin')

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        data = {
            'X': torch.from_numpy(self.X[idx]),
            'Y': torch.from_numpy(self.Y[idx]),
            'P': torch.from_numpy(self.P[idx])
        }
        return data


class PFNNBinaryDataset(Dataset):
    def __init__(self,
                 base_path = 'dataset',
                 name = 'database.npz'):
        self.base_path = base_path
        database_path = os.path.join(base_path, name)
        print("Loading database...")
        database = np.load(database_path)
        self.X = database['Xun'].astype(np.float32)
        self.Y = database['Yun'].astype(np.float32)
        self.P = database['Pun'].astype(np.float32)[..., None]
        self.calculate_statistics()
        print("Load completed")
        print(f"Input shape: {self.X.shape}")
        print(f"Output shape: {self.Y.shape}")
        print(f"Phase shape: {self.P.shape}")
        self.length = len(self.X)

    def __len__(self):
        return self.length

    def __getitem__(self, idx):
        data = {
            'X': torch.from_numpy(self.X[idx]),
            'Y': torch.from_numpy(self.Y[idx]),
            'P': torch.from_numpy(self.P[idx])
        }
        return data

    def calculate_statistics(self):
        """ Calculate Mean and Std """
        Xmean, Xstd = self.X.mean(axis=0), self.X.std(axis=0)
        Ymean, Ystd = self.Y.mean(axis=0), self.Y.std(axis=0)

        j = 31
        w = ((60 * 2) // 10)

        Xstd[w * 0:w * 1] = Xstd[w * 0:w * 1].mean()  # Trajectory Past Positions
        Xstd[w * 1:w * 2] = Xstd[w * 1:w * 2].mean()  # Trajectory Future Positions
        Xstd[w * 2:w * 3] = Xstd[w * 2:w * 3].mean()  # Trajectory Past Directions
        Xstd[w * 3:w * 4] = Xstd[w * 3:w * 4].mean()  # Trajectory Future Directions
        Xstd[w * 4:w * 10] = Xstd[w * 4:w * 10].mean()  # Trajectory Gait

        """ Mask Out Unused Joints in Input """

        joint_weights = np.array([
            1,
            1e-10, 1, 1, 1, 1,
            1e-10, 1, 1, 1, 1,
            1e-10, 1, 1,
            1e-10, 1, 1,
            1e-10, 1, 1, 1, 1e-10, 1e-10, 1e-10,
            1e-10, 1, 1, 1, 1e-10, 1e-10, 1e-10]).repeat(3)

        Xstd[w * 10 + j * 3 * 0:w * 10 + j * 3 * 1] = Xstd[w * 10 + j * 3 * 0:w * 10 + j * 3 * 1].mean() / (
                    joint_weights * 0.1)  # Pos
        Xstd[w * 10 + j * 3 * 1:w * 10 + j * 3 * 2] = Xstd[w * 10 + j * 3 * 1:w * 10 + j * 3 * 2].mean() / (
                    joint_weights * 0.1)  # Vel
        Xstd[w * 10 + j * 3 * 2:] = Xstd[w * 10 + j * 3 * 2:].mean()  # Terrain

        Ystd[0:2] = Ystd[0:2].mean()  # Translational Velocity
        Ystd[2:3] = Ystd[2:3].mean()  # Rotational Velocity
        Ystd[3:4] = Ystd[3:4].mean()  # Change in Phase
        Ystd[4:8] = Ystd[4:8].mean()  # Contacts

        Ystd[8 + w * 0:8 + w * 1] = Ystd[8 + w * 0:8 + w * 1].mean()  # Trajectory Future Positions
        Ystd[8 + w * 1:8 + w * 2] = Ystd[8 + w * 1:8 + w * 2].mean()  # Trajectory Future Directions

        Ystd[8 + w * 2 + j * 3 * 0:8 + w * 2 + j * 3 * 1] = Ystd[
                                                            8 + w * 2 + j * 3 * 0:8 + w * 2 + j * 3 * 1].mean()  # Pos
        Ystd[8 + w * 2 + j * 3 * 1:8 + w * 2 + j * 3 * 2] = Ystd[
                                                            8 + w * 2 + j * 3 * 1:8 + w * 2 + j * 3 * 2].mean()  # Vel
        Ystd[8 + w * 2 + j * 3 * 2:8 + w * 2 + j * 3 * 3] = Ystd[
                                                            8 + w * 2 + j * 3 * 2:8 + w * 2 + j * 3 * 3].mean()  # Rot

        """ Save Mean / Std / Min / Max """

        Xmean.astype(np.float32).tofile(os.path.join(self.base_path, 'Xmean.bin'))
        Ymean.astype(np.float32).tofile(os.path.join(self.base_path, 'Ymean.bin'))
        Xstd.astype(np.float32).tofile(os.path.join(self.base_path, 'Xstd.bin'))
        Ystd.astype(np.float32).tofile(os.path.join(self.base_path, 'Ystd.bin'))
        self.X = (self.X - Xmean) / Xstd
        self.Y = (self.Y - Ymean) / Ystd

if __name__ == "__main__":
    dset = PFNNBinaryDataset()
    dloader = DataLoader(dset, batch_size=32, shuffle=False)
    for batch in dloader:
        for k,v in batch.items():
            print(f'{k} shape: {v.shape}')
